import requests
import json
from typing import Optional, Dict, Any


class Text2SQLDemo:
    """
    Text2SQL 转换演示类
    """

    def __init__(self, api_key: str, workflow_id: str, base_url: str = "http://localhost"):
        """
        初始化Text2SQL演示

        Args:
            api_key (str): Dify API密钥
            workflow_id (str): 工作流ID
            base_url (str): API基础URL
        """
        self.api_key = api_key
        self.workflow_id = workflow_id
        self.base_url = base_url.rstrip('/')

    def text_to_sql(self, natural_query: str, user_id: str = "default_user") -> Optional[Dict[Any, Any]]:
        """
        将自然语言转换为SQL查询

        Args:
            natural_query (str): 自然语言查询
            user_id (str): 用户标识

        Returns:
            dict: 包含SQL查询结果的字典
        """
        url = f"{self.base_url}/v1/workflows/{self.workflow_id}/run"

        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json"
        }

        payload = {
            "inputs": {
                "query": natural_query
            },
            "response_mode": "blocking",
            "user": user_id
        }

        try:
            response = requests.post(url, headers=headers, json=payload, timeout=30)
            response.raise_for_status()
            return response.json()
        except requests.exceptions.RequestException as e:
            print(f"Text2SQL转换失败: {e}")
            return None

    def execute_sql_query(self, sql_query: str, database_config: Dict[str, Any]) -> Optional[Dict[Any, Any]]:
        """
        执行SQL查询（示例方法，需要根据实际数据库连接实现）

        Args:
            sql_query (str): SQL查询语句
            database_config (dict): 数据库连接配置

        Returns:
            dict: 查询结果
        """
        # 这里应该实现实际的数据库查询逻辑
        # 例如使用 sqlite3, pymysql, psycopg2 等
        print(f"执行SQL查询: {sql_query}")
        print(f"数据库配置: {database_config}")

        # 模拟返回结果
        return {
            "status": "success",
            "data": [
                {"id": 1, "name": "示例数据", "value": 100}
            ],
            "row_count": 1
        }


def demo_usage():
    """
    演示如何使用Text2SQL功能
    """
    # 配置参数
    API_KEY = "your-dify-api-key"
    WORKFLOW_ID = "your-text2sql-workflow-id"

    # 创建Text2SQL演示实例
    text2sql = Text2SQLDemo(API_KEY, WORKFLOW_ID)

    # 自然语言查询示例
    queries = [
        "查询所有用户的姓名和邮箱",
        "找出销售额大于10000的订单",
        "统计每个部门的员工数量",
        "获取最近30天的登录记录"
    ]

    # 执行转换
    for query in queries:
        print(f"\n{'=' * 50}")
        print(f"自然语言查询: {query}")
        print(f"{'=' * 50}")

        # 调用Text2SQL转换
        result = text2sql.text_to_sql(query)

        if result:
            print("转换结果:")
            print(json.dumps(result, indent=2, ensure_ascii=False))

            # 如果返回了SQL查询语句，可以进一步执行
            if 'data' in result and 'sql' in result['data']:
                sql_query = result['data']['sql']
                print(f"\n生成的SQL: {sql_query}")

                # 执行SQL查询（示例）
                db_config = {
                    "host": "localhost",
                    "port": 3306,
                    "database": "test_db",
                    "username": "user",
                    "password": "password"
                }

                query_result = text2sql.execute_sql_query(sql_query, db_config)
                if query_result:
                    print("查询执行结果:")
                    print(json.dumps(query_result, indent=2, ensure_ascii=False))
        else:
            print("转换失败")


if __name__ == "__main__":
    demo_usage()
